"""
@Type doc
@Author xjp
@CreateDate 2025-01-17_09:30:29
@Description Llama模型
@Version v1.0
@Copyright Copyright (c) 2025 by xiaoxiang ,ALL Rights Reserved
"""

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# 利用代码下载模型文件
#from modelscope import snapshot_download
#model_id = snapshot_download("LLM-Research/Meta-Llama-3-8B-Instruct")

"""
@Type class
@Author xjp
@CreateDate 2025-01-17_09:30:58
@Description 调用Llama3模型的类 
"""
class Llama3:
    def __init__(self):
        # 使用transformer加载模型

        # 使用的设备
        # device ="cuda"  # 使用GPU计算
        device = "cpu" # 使用CPU计算

        # 模型文件路径
        modelDir = "./model/Llama3"

        # 加载了一个因果语言模型。
        # pretrained_model_name_or_path 是模型文件所在的目录。# torch_dtype="auto" 自动选择最优的数据类型以平衡性能和精度。# device_map="auto" 自动将模型的不同部分映射到可用的设备上。
        self.model= AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=modelDir,torch_dtype='auto',device_map="auto")

        
        # 加载与模型相匹配的分词器。分词器用于将文本转换成模型能够理解和处
        self.tokenizer=AutoTokenizer.from_pretrained(modelDir)
    
    """
    @Type function
    @Author xjp
    @CreateDate 2025-01-17_10:06:48
    @Description 分词工具 
    @Param 
    @Return 
    """
    def __wordSegmentation(self, prompt:str):
        #加载与模型相匹配的分词器。分词器用于将文本转换成模型能够理解和处理
        messages=[{'role':'system','content':'You are a helpful assistant system'},{'role': 'user','content': prompt}]
        # 使用分词器的 apply_chat_template 方法将上面定义的消息列表转护# tokenize=False 表示此时不进行令牌化，add_generation_promp
        text =self.tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True)

        #将处理后的文本令牌化并转换为模型输入张量，然后将这些张量移至之前
        self.model_inputs=self.tokenizer([text],return_tensors="pt").to('cuda')
        self.input_ids = self.tokenizer.encode(text, return_tensors='pt')
        self.attention_mask=torch.ones(self.input_ids.shape,dtype=torch.long,device="cuda")
        
    

    """
    @Type function
    @Author xjp
    @CreateDate 2025-01-17_10:11:03
    @Description 运行模型 
    @Param 
    @Return 
    """
    def __runModel(self):
        generated_ids = self.model.generate(
	    self.model_inputs.input_ids,
                      attention_mask=self.attention_mask,
	    max_new_tokens=512,
                     pad_token_id=self.tokenizer.eos_token_id
	    )
	
	    # 对输出进行解码
        response=self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)

        return response
    
    """
    @Type function
    @Author xjp
    @CreateDate 2025-01-17_10:12:41
    @Description 运行模型 
    @Param 
    @Return 
    """
    def run(self, prompt:str):
        self.__wordSegmentation(prompt)
        res = self.__runModel()
        return res



if __name__ == "__main__":
    model = Llama3()
    print("111")
    res = model.run("你好呀,请介绍一下自己吧!!!")
    print("222")
    print(res)


